/**
 * 
 */
package datastructure;

import java.util.PriorityQueue;
import java.util.List;
import java.util.ArrayList;

/**
 * @author michaelyi
 *
 */
public class KWayMerge {

	static class Pair implements Comparable<Pair>{
		List<Integer> list;
		int index;	// current position in the list
		Pair(List<Integer> list, int index) {
			assert(list != null);
			assert(index < list.size());
			this.list = list;
			this.index = index;	
		}
		
		@Override
		public int compareTo(Pair p) {
			int result = 1;
			if (p == null) return result;
			
			if (this.list.get(index) < p.list.get(p.index))
				result = -1;
			else if (this.list.get(index) == p.list.get(p.index))
				result = 0;
			
			return result;
		}
	}
	
	/**
	 * 
	 */
	public KWayMerge() {
		// TODO Auto-generated constructor stub
	}

	public static List<Integer> merge2SortedList(List<Integer> list1, List<Integer> list2) {
		if (list1 == null && list2 == null) return null;
		if (list1 == null) return new ArrayList<Integer>(list2);
		if (list2 == null) return new ArrayList<Integer>(list1);
		
		List<Integer> result = new ArrayList<Integer>();
		int i1 = 0, i2 = 0;
		while (i1 < list1.size() && i2 < list2.size()) {
			if (list1.get(i1) < list2.get(i2))	result.add(list1.get(i1++));
			else result.add(list2.get(i2++));
		}
		
		while (i1 < list1.size()) {
			result.add(list1.get(i1++));
		}
		
		while (i2 < list2.size()) {
			result.add(list2.get(i2++));
		}
		return result;
	}
	
	public static List<Integer> mergeKsortedListPQ(List< List<Integer> > list) {
		if (list == null) return null;
		
		List <Integer> result = new ArrayList<Integer>();

		ArrayList<Pair> al = new ArrayList<Pair>(list.size());
		for (List<Integer> l : list) {
			al.add(new Pair(l, 0));
		}
		
		PriorityQueue<Pair> pq = new PriorityQueue<Pair>(al);
		
		while (pq.size() > 0) {
			result.add(pq.peek().list.get(pq.peek().index));
			if (pq.peek().index == pq.peek().list.size()-1) {
				pq.poll();
			}
			else {
				pq.peek().index++;
				pq.add(pq.poll());
			}
		}
		
		return result;
	}
	
	public static List<Integer> mergeKsortedList(List< List<Integer> > list) {
		if (list == null) return null;
		
		List <Integer> result = new ArrayList<Integer>();
		
		Pair[] heap = new Pair[list.size()];
		for (int i = 0; i < list.size(); ++i) {
			heap[i] = new Pair(list.get(i), 0);
		}
		

		int heapSize = heap.length;
		heapify(heap, heapSize);
		while (heapSize > 0) {
			result.add(heap[0].list.get(heap[0].index));
			if (heap[0].index == heap[0].list.size()-1) {
				heap[0] = heap[--heapSize];
			}
			else {
				heap[0].index++;
			}
			if (heapSize > 1)
				shiftDown(heap, heapSize, 0);
		}
		
		return result;
	}
	
	private static void shiftDown(Pair[] heap, int heapSize, int index) {
		assert(heapSize <= heap.length);
		assert(index >= 0 && index < heapSize);
		
		int nextIndex = index * 2 + 1;
		while (nextIndex < heapSize) {
			int swapIndex = index;
			if (heap[swapIndex].list.get(heap[swapIndex].index) > heap[nextIndex].list.get(heap[nextIndex].index)) {
				swapIndex = nextIndex;
			}
			if (++nextIndex < heapSize && heap[swapIndex].list.get(heap[swapIndex].index) 
				> heap[nextIndex].list.get(heap[nextIndex].index)) {
				swapIndex = nextIndex;
			}
			
			if (swapIndex == index) {
				break;
			}

			swap(heap, index, swapIndex);
			index = swapIndex;
			nextIndex = index * 2 + 1;
		}
	}
	
	private static void swap(Pair[] heap, int index1, int index2) {
		Pair temp = heap[index1];
		heap[index1] = heap[index2];
		heap[index2] = temp;
	}
	
	private static void heapify(Pair[] heap, int heapSize) {
		assert(heapSize <= heap.length);
		
		for (int i = heapSize / 2 - 1; i >= 0; --i) {
			shiftDown(heap, heapSize, i);
		}
	}
	

	
	/**
	 * @param args
	 */
	public static void main(String[] args) {
		// TODO Auto-generated method stub
		System.out.printf("2 Way Merge");
		List<List<Integer>> list = new ArrayList<List<Integer>>();
		list.add(new ArrayList<Integer>());
		list.add(new ArrayList<Integer>());
		list.add(new ArrayList<Integer>());
		list.get(0).add(2); list.get(0).add(3);	list.get(0).add(4); list.get(0).add(5); list.get(0).add(7);
		list.get(1).add(0);	list.get(1).add(4); list.get(1).add(5); list.get(1).add(6); list.get(1).add(8);
		list.get(2).add(1);	list.get(2).add(2); list.get(2).add(3); list.get(2).add(7); list.get(2).add(9);
		
		List<Integer> result = merge2SortedList(list.get(0), list.get(1));
		for (int i : result) {
			System.out.printf("%d, ", i);
		}
		
		System.out.println();
		
		result = mergeKsortedList(list);
		//result = mergeKsortedListPQ(list);
		for (int i : result) {
			System.out.printf("%d, ", i);
		}

	}

}


